package com.tomshidi.security.distributed.uaa.redis;

import org.springframework.data.redis.core.RedisTemplate;
import org.springframework.data.redis.core.StringRedisTemplate;
import org.springframework.security.crypto.keygen.Base64StringKeyGenerator;
import org.springframework.security.oauth2.common.exceptions.InvalidGrantException;
import org.springframework.security.oauth2.common.util.SerializationUtils;
import org.springframework.security.oauth2.provider.OAuth2Authentication;
import org.springframework.security.oauth2.provider.code.AuthorizationCodeServices;
import org.springframework.security.oauth2.provider.code.RandomValueAuthorizationCodeServices;
import org.springframework.web.context.request.RequestAttributes;
import org.springframework.web.context.request.RequestContextHolder;
import org.springframework.web.context.request.ServletRequestAttributes;

import javax.servlet.http.HttpServletRequest;
import java.util.Base64;
import java.util.Set;
import java.util.stream.Collectors;

/**
 * Redis实现的授权码存储服务 copy from {@link RandomValueAuthorizationCodeServices}
 * <p>
 * {@link #createAuthorizationCode} 方法将请求的sessionId放到details中，以便{@link JwtUtil#getSessionAccessToken}可以获取
 * 此session所生成的token，用于单点退出
 * </p>
 *
 */

public class RedisAuthorizationCodeServices implements AuthorizationCodeServices {

    private RedisTemplate redisTemplate;


    public RedisAuthorizationCodeServices(RedisTemplate redisTemplate) {
        this.redisTemplate = redisTemplate;
    }

    @Override
    public String createAuthorizationCode(OAuth2Authentication authentication) {
        String code = new Base64StringKeyGenerator(Base64.getUrlEncoder()).generateKey();
        HttpServletRequest request = getRequest();

        if (request != null) {
            String sessionId = request.getSession().getId();
            redisTemplate.opsForValue().set("base:oauth2:code-session:test:" + code,sessionId,60L);
        }
        this.store(code, authentication);
        return code;
    }

    private HttpServletRequest getRequest() {
        RequestAttributes requestAttributes = RequestContextHolder.getRequestAttributes();
        return requestAttributes == null ? null : ((ServletRequestAttributes)requestAttributes).getRequest();
    }

    @Override
    public OAuth2Authentication consumeAuthorizationCode(String code) throws InvalidGrantException {
        OAuth2Authentication auth = this.remove(code);
        if (auth == null) {
            throw new InvalidGrantException("Invalid authorization code: " + code);
        }
        return auth;
    }

    private void store(String code, OAuth2Authentication authentication) {
        redisTemplate.opsForValue().set("base:oauth2:code:test:" + code,SerializationUtils.serialize(authentication),60L);
    }

    private OAuth2Authentication remove(String code) {

        byte[] valueBytes = (byte[]) redisTemplate.opsForValue().get("base:oauth2:code:test:" + code);

        if (valueBytes != null) {
            OAuth2Authentication authentication = SerializationUtils.deserialize(valueBytes);
            redisTemplate.delete("base:oauth2:code:test:" + code);
            return authentication;
        } else {
            return null;
        }
    }


}
